import csv
import re
import pandas as pd
from langchain_core.prompts import ChatPromptTemplate
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import CountVectorizer

# 假设TOKENIZER和LLM已经定义
# from utils.instances import LLM, TOKENIZER
# import utils.configFinRAG as configFinRAG
# from utils import prompts

vectorizer = CountVectorizer()

def tokenize_and_filter(text, pattern):
    filtered_text = re.sub(pattern, ' ', text)
    return TOKENIZER(filtered_text)['input_ids']

def calculate_similarity(new_token, example_vectors):
    new_vector = vectorizer.transform([' '.join(map(str, new_token))])
    return cosine_similarity(new_vector, example_vectors).flatten()

def generate_answer(question, fa, llm, example_questions, example_infos, example_fas, example_tokens, example_vectors, example_num=5):
    pattern = r'\d{8}'
    question_token = tokenize_and_filter(question, pattern)
    similarity = calculate_similarity(question_token, example_vectors)

    top_indices = similarity.argsort()[-example_num:][::-1]
    prompt = ChatPromptTemplate.from_template(prompts.ANSWER_TEMPLATE)
    examples = '\n'.join(
        f"问题：{example_questions[i]}\n资料：{example_infos[i]}\n答案：{example_fas[i]}\n"
        for i in top_indices
    )

    chain = prompt | llm
    response = chain.invoke({"examples": examples, "FA": fa, "question": question})
    return response.content

def main():
    # 读取问题和FA模板
    sql_examples_file = pd.read_csv(configFinRAG.sql_examples_path)
    example_questions = sql_examples_file['问题'].tolist()
    example_infos = sql_examples_file['资料'].tolist()
    example_fas = sql_examples_file['FA'].tolist()

    # Tokenize所有问题并计算向量
    example_tokens = [tokenize_and_filter(q, r'\d{8}') for q in example_questions]
    example_text_vectors = [' '.join(map(str, tokens)) for tokens in example_tokens]
    vectorizer.fit(example_text_vectors)
    example_vectors = vectorizer.transform(example_text_vectors)

    # 读取待回答的问题
    result_csv_file = pd.read_csv(configFinRAG.question_sql_check_path)

    # 写入答案
    with open(configFinRAG.answer_path, 'w', newline='', encoding='utf-8-sig') as answer_file:
        csvwriter = csv.writer(answer_file)
        csvwriter.writerow(['问题id', '问题', '资料', 'FA'])

        for _, row in result_csv_file.iterrows():
            if row['flag']:
                result = generate_answer(row['问题'], row['执行结果'], LLM,
                                         example_questions, example_infos, example_fas, example_tokens, example_vectors)
                csvwriter.writerow([row['问题id'], row['问题'], row['执行结果'], result])

if __name__ == '__main__':
    main()
